import geopandas as gpd
import numpy as np
import json
from pykrige.ok import OrdinaryKriging
import pandas as pd
import matplotlib.pyplot as plt
import plotnine
from plotnine import *
from shapely import Point
import rasterio
from rasterio.mask import mask
from rasterio.transform import from_bounds
from shapely.geometry import mapping

def crop_with_geojson(data, region, transform):
    import numpy as np
    with rasterio.io.MemoryFile() as memfile:
        with memfile.open(
            driver='GTiff',
            height=data.shape[0],
            width=data.shape[1],
            count=1,
            crs='EPSG:4326',
            dtype='float32',
            transform=transform
        ) as dataset:
            dataset.write(data, 1)
            out_image, _ = mask(dataset, [mapping(region)], crop=False, all_touched=True)
            out_image[out_image == 0] = np.nan
            return out_image[0]
# 读取 GeoJSON 文件（可本地或远程）
gdf = gpd.read_file("100000.geo.json")
# gdf = gdf.to_crs(epsg=4326)
gdf.set_crs("EPSG:4326")
# 查看地图数据的基本信息
print(gdf.crs)  # 坐标参考系统
print(gdf.columns)  # 字段名
print(gdf.bounds)
gdf_bounds = gdf.total_bounds
grid_lon = np.linspace(gdf_bounds[0], gdf_bounds[2], 400)
grid_lat = np.linspace(gdf_bounds[1], gdf_bounds[3], 400)

with open("data2.json", "r") as f:
    data_list = json.load(f)
# data = np.around(np.array(data_list), 6)
data = np.array(data_list)
points_lat = data[:, 0]
points_lon = data[:, 1]
points_aqi = data[:, 2]

OK = OrdinaryKriging(points_lon, points_lat, points_aqi, variogram_model='gaussian', nlags=6)
z1, ss1 = OK.execute('grid', grid_lon, grid_lat)
z1 = np.flipud(z1)
# 转换成网格
xgrid, ygrid = np.meshgrid(grid_lon, grid_lat)
df_grid = pd.DataFrame({
    'long': xgrid.flatten(),
    'lat': ygrid.flatten(),
    'Krig_gaussian': z1.flatten()
})

# 转换为GeoDataFrame
geometry = [Point(xy) for xy in zip(df_grid['long'], df_grid['lat'])]
print(gdf.crs)
grid_gdf = gpd.GeoDataFrame(df_grid, geometry=geometry,crs="EPSG:4326")
print(gdf.crs)
# 进行空间裁剪
grid_gdf = gpd.sjoin(grid_gdf, gdf, how='inner', predicate='within')

# 将插值网格数据整理
# df_grid = pd.DataFrame(dict(long=xgrid.flatten(), lat=ygrid.flatten()))

# 添加插值结果
# df_grid["Krig_gaussian"] = z1.flatten()

# plotnine.options.figure_size = (5, 4.5)
Krig_inter_grid = (ggplot() +
                   geom_tile(grid_gdf, aes(x='long', y='lat', fill='Krig_gaussian'), size=0.1) +
                   geom_map(gdf, fill='none', color='none') +
                   scale_fill_gradientn(colors=['#9cd84e', '#facf39', '#f99049',
                                                '#f65e5f', '#a070b6', '#a06a7b'],
                                        breaks=[50, 100, 150, 200, 300, 400],
                                        name='AQI') +
                   # scale_fill_cmap(cmap_name='Spectral_r', name='Krig_gaussian_result',
                   #                 breaks=[50, 100, 150, 200, 300, 400]
                   #                 ) +
                   # scale_x_continuous(breaks=[117, 118, 119, 120, 121, 122]) +
                   theme(
                       text=element_text(family="Consolas", size=12),
                       # 修改背景
                       legend_position='none',  # 隐藏图例
                       axis_title=element_blank(),  # 去掉坐标轴标题
                       axis_text=element_blank(),  # 去掉坐标轴刻度文字
                       axis_ticks=element_blank(),  # 去掉坐标轴刻度线
                       panel_background=element_blank(),
                       axis_ticks_major_x=element_blank(),
                       axis_ticks_major_y=element_blank(),
                       panel_grid_major=element_blank(),  # 去掉网格线
                       panel_grid_minor=element_blank()
                   ))

# Krig_inter_grid.save("img-2025-01-10-16.png", dpi=300, bbox_inches='tight',  # 紧凑布局
#                      transparent=True,  # 透明背景
#                      pad_inches=0)


region = gdf.geometry.unary_union
# 计算 transform
transform = from_bounds(
    gdf_bounds[0], gdf_bounds[1], gdf_bounds[2], gdf_bounds[3],
    z1.shape[1], z1.shape[0]
)
# GeoJSON 区域裁剪
z1_crop = crop_with_geojson(z1, region, transform)
# 翻转
z1_crop = np.flipud(z1_crop)



# 用 matplotlib 保存图片，确保坐标与 imageExtent 一致
plt.figure(figsize=(8, 6), dpi=300)
plt.axis('off')
plt.imshow(
    z1_crop,
    extent=[grid_lon.min(), grid_lon.max(), grid_lat.min(), grid_lat.max()],
    origin='lower',  # 保证左下为原点
    cmap=plt.cm.get_cmap('Spectral_r')
)
plt.savefig('img-2025-01-10-16.png',
            bbox_inches='tight', pad_inches=0, transparent=True)
plt.close()